#######################################################################
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# Openfuyao is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#     http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
# PURPOSE.
# See the Mulan PSL v2 for more details.
#######################################################################

#!/bin/bash

# 获取当前脚本所在目录的绝对路径
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)

# 默认参数
RUN_TIMES=1
TEST_TYPE="default"
BASE_YAML="${SCRIPT_DIR}/test-job-wrk-client.yaml"
NAMESPACE="default"
JOB_PREFIX="wrk-client-test"
LOG_DIR="/usr/local/nginx-test"
CPU_REQUEST="2"
CPU_LIMIT="2"
CONCURRENCY="200"
DURATION="300s"
URL=""

# 将时间字符串转换为秒数
function convert_to_seconds() {
    local time_str=$1
    local total_seconds=0
    
    # 提取分钟部分
    if [[ $time_str =~ ([0-9]+)m ]]; then
        total_seconds=$((total_seconds + ${BASH_REMATCH[1]} * 60))
    fi
    
    # 提取秒数部分
    if [[ $time_str =~ ([0-9]+)s ]]; then
        total_seconds=$((total_seconds + ${BASH_REMATCH[1]}))
    fi
    
    echo $total_seconds
}

# 帮助信息
function show_help() {
    echo "用法: $0 -c <并发数> -d <持续时间> -u <URL> [选项]"
    echo "必选参数:"
    echo "  -u <URL>       测试目标URL (必填)"
    echo "可选参数:"
    echo "  -t <类型>       测试类型，用于日志文件名 (默认: default)"
    echo "  -c <并发数>     并发连接数 (默认: 200)"
    echo "  -d <持续时间>   压测持续时间 (如: 300s, 5m, 1m30s) (默认: 300s)"
    echo "  -n <次数>       循环运行测试的次数 (默认: 1)"
    echo "  -y <yaml>      指定基础 YAML 文件路径 (默认: test-job-wrk-client.yaml)"
    echo "  -req <CPU>      CPU requests (例如: 2, 1.5, 500m) (默认: 2)"
    echo "  -limit <CPU>    CPU limits (例如: 2, 1.5, 500m) (默认: 2)"
    echo "  -h             显示此帮助信息"
    exit 1
}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -c)
            CONCURRENCY="$2"
            shift 2
            ;;
        -d)
            DURATION="$2"
            shift 2
            ;;
        -t)
            TEST_TYPE="$2"
            shift 2
            ;;
        -n)
            RUN_TIMES="$2"
            shift 2
            ;;
        -req)
            CPU_REQUEST="$2"
            shift 2
            ;;
        -limit)
            CPU_LIMIT="$2"
            shift 2
            ;;
        -y)
            BASE_YAML="$2"
            shift 2
            ;;
        -u)
            URL="$2"
            shift 2
            ;;
        -h)
            show_help
            ;;
        *)
            echo "未知参数: $1" >&2
            show_help
            ;;
    esac
done

# 验证参数
if [[ -z "$URL" ]]; then
    echo "错误: URL参数(-u)是必填项" >&2
    show_help
fi

if ! [[ "$CONCURRENCY" =~ ^[0-9]+$ ]]; then
    echo "错误: 并发数必须是数字" >&2
    show_help
fi

if ! [[ "$DURATION" =~ ^([0-9]+[sm]?)+$ ]]; then
    echo "错误: 持续时间格式不正确 (示例: 120s, 5m, 1m30s)" >&2
    show_help
fi

if ! [[ "$RUN_TIMES" =~ ^[0-9]+$ ]]; then
    echo "错误: 运行次数必须是数字" >&2
    show_help
fi

if [[ -n "$CPU_REQUEST" && ! "$CPU_REQUEST" =~ ^([0-9]+m|[0-9]+(\.[0-9]+)?)$ ]]; then
    echo "错误: CPU requests 格式不正确 (示例: 2, 1.5, 500m)" >&2
    show_help
fi

if [[ -n "$CPU_LIMIT" && ! "$CPU_LIMIT" =~ ^([0-9]+m|[0-9]+(\.[0-9]+)?)$ ]]; then
    echo "错误: CPU limits 格式不正确 (示例: 2, 1.5, 500m)" >&2
    show_help
fi

# 创建日志目录
mkdir -p "$LOG_DIR"

# 将持续时间转换为秒数
DURATION_SECONDS=$(convert_to_seconds "$DURATION")

# 循环运行测试
for i in $(seq 1 $RUN_TIMES); do
    # 生成日志文件名 格式: log-{type}-{concurrency}-{duration_seconds}-{index}
    LOG_FILE="log-${TEST_TYPE}-t2-c${CONCURRENCY}-d${DURATION_SECONDS}-${i}"
    JOB_NAME="${JOB_PREFIX}-${TEST_TYPE}-${CONCURRENCY}-${i}"

    # 创建临时 YAML 文件
    TEMP_YAML="${SCRIPT_DIR}/wrk-client-${TEST_TYPE}-${CONCURRENCY}-${i}.yaml"
    cp "$BASE_YAML" "$TEMP_YAML"

    # 使用yq修改YAML文件参数
    yq -i -y ".metadata.name = \"${JOB_NAME}\"" "$TEMP_YAML"
    
    # 使用yq修改环境变量
    yq -i -y '(.spec.template.spec.containers[0].env[] | select(.name == "CONCURRENCY").value) = "'"${CONCURRENCY}"'"' "$TEMP_YAML"
    yq -i -y '(.spec.template.spec.containers[0].env[] | select(.name == "DURATION").value) = "'"${DURATION_SECONDS}s"'"' "$TEMP_YAML"
    yq -i -y '(.spec.template.spec.containers[0].env[] | select(.name == "URL").value) = "'"${URL}"'"' "$TEMP_YAML"
    yq -i -y '(.spec.template.spec.containers[0].env[] | select(.name == "RESULT_FILE").value) = "'"${LOG_FILE}"'"' "$TEMP_YAML"

    # 使用yq修改CPU requests和limits
    yq -i -y ".spec.template.spec.containers[0].resources.requests.cpu = \"${CPU_REQUEST}\"" "$TEMP_YAML"
    yq -i -y ".spec.template.spec.containers[0].resources.limits.cpu = \"${CPU_LIMIT}\"" "$TEMP_YAML"

    # 应用 Job
    echo "创建 Job ${JOB_NAME} (类型: ${TEST_TYPE}, 并发: ${CONCURRENCY}, 持续时间: ${DURATION} (${DURATION_SECONDS}s), CPU: ${CPU_REQUEST}/${CPU_LIMIT})..."    
    kubectl apply -f "$TEMP_YAML" -n "$NAMESPACE"

    # 等待 Job 完成
    echo "等待 Job ${JOB_NAME} 完成..."
    while true; do
        STATUS=$(kubectl get job "$JOB_NAME" -n "$NAMESPACE" -o jsonpath='{.status.succeeded}')
        if [[ "$STATUS" == "1" ]]; then
            echo "Job ${JOB_NAME} 已成功完成!"
            break
        fi
        
        # 检查是否失败
        FAILED=$(kubectl get job "$JOB_NAME" -n "$NAMESPACE" -o jsonpath='{.status.failed}')
        if [[ -n "$FAILED" && "$FAILED" -gt 0 ]]; then
            echo "Job ${JOB_NAME} 失败!"
            kubectl describe job "$JOB_NAME" -n "$NAMESPACE"
            kubectl logs job/"$JOB_NAME" -n "$NAMESPACE"
            exit 1
        fi
        
        sleep 10
    done

    # 清理资源
    kubectl delete job "$JOB_NAME" -n "$NAMESPACE"
    rm -f "$TEMP_YAML"

    echo "第 $i 次测试完成! 结果保存在: ${LOG_FILE}"
    echo "------------------------"
    sleep 5  # 测试间隔
done

echo "所有测试已完成!"